% For a DSGE model approximated up to second order by projection, this function calculates 
% the unconditional first and second moments based on the the pruning
% scheme. The system reads
%     xf(t+1) = h0 + hx*xf(t) + eta*eps(t);
%     xs(t+1) = hx*xs(t) + reshape(hxx,nx,nx^2)*kron(xf(t),xf(t));
%     y(t+1)  = g0 + gx*(xf(t+1)+xs(t+1)) + reshape(gxx,ny,nx^2)*kron(xf(t+1),xf(t+1))
%
% IMPORTANT: Here the Lyaponov method is used to solve for the variances
%
% INPUTS:
% g0,gx,... : The second order approximation of the DSGE model by projection
% vectorMom3 : vector with dimensions 1*ne with third moments of shocks
% vectorMom4 : vector with dimensions 1*ne with fourth moments of shocks
% auto_lag   : the lag lenght for the auto-covariances (they need to be divided
%              by the standard deviations to get the correlations)
%
% OUTPUTS: 
% See below
% By M. Andreasen, 2022

function out = UnconditionalMomentsProj_2nd_Lyap(g0,gx,gxx,h0,hx,hxx,eta,...
    vectorMom3,vectorMom4,auto_lag,Var_zold)

% The dimensions
ny = length(g0);
[nx,ne] = size(eta);

% The auxiliary system z=[xf xs kron(xf,xf)]
% The loading matrix
A = [hx                        ,zeros(nx,nx)       ,zeros(nx,nx^2);
    zeros(nx,nx)               ,hx                 ,reshape(hxx,nx,nx*nx);
    kron(h0,hx)+kron(hx,h0)    ,zeros(nx^2,nx)     ,kron(hx,hx)          ];
% The constant term
c = [h0;
    zeros(nx,1);
    kron(h0,h0)+kron(eta,eta)*reshape(eye(ne),ne^2,1)];

% the mean value of the states in the auxiliary system
%E_z    = inv(eye(2*nx+nx^2)-A)*c;
E_z     = (eye(2*nx+nx^2)-A)\c;
E_xf    = E_z(1:nx,1);
E_xs    = E_z(nx+1:2*nx,1);
E_xf_xf = reshape(E_z(2*nx+1:2*nx+nx^2,1),nx,nx);

% The B matrix 
B = zeros(2*nx+nx^2,ne+ne^2+2*ne*nx);
B(1:nx,1:ne)    = eta;
B(2*nx+1:end,:) = [kron(h0,eta)+kron(eta,h0) kron(eta,eta) kron(eta,hx) kron(hx,eta)];

% The size of the variance of the innovations, VarZi
%VarZi = zeros(ne+ne^2+2*ne*nx,ne+ne^2+2*ne*nx);

E_eps2_eps2 = zeros(ne^2,ne^2);
index1 = 0;
for phi4=1:ne
    for phi1=1:ne
        index1 = index1 + 1;
        index2 = 0;
        for phi3=1:ne
            for phi2=1:ne
                index2 = index2 + 1;
                % The second moments
                if (phi1 == phi2 && phi3 == phi4 && phi1 ~= phi4)
                    E_eps2_eps2(index1,index2) = 1;
                elseif (phi1 == phi3 && phi2 == phi4 && phi1 ~= phi2)
                    E_eps2_eps2(index1,index2) = 1;
                elseif (phi1 == phi4 && phi2 == phi3 && phi1 ~= phi2)                    
                    E_eps2_eps2(index1,index2) = 1;
                elseif (phi1 == phi2 && phi1 == phi3 && phi1 == phi4)
                    % The fourth moments
                    E_eps2_eps2(index1,index2) = vectorMom4(1,phi1);
                end
            end
        end
    end
end                  

E_xfeps_epsxf = zeros(nx*ne,nx*ne);
index1 = 0;
for gama1=1:nx
    for phi1=1:ne
        index1 = index1 + 1;
        index2 = 0;
        for phi2=1:ne
            for gama2=1:nx
                index2 = index2 + 1;
                if phi1 == phi2 
                    E_xfeps_epsxf(index1,index2) = E_xf_xf(gama1,gama2);
                end
            end
        end
    end
end
E_epsxf_xfeps = E_xfeps_epsxf';

vecI_vecI = reshape(eye(ne),ne^2,1)*reshape(eye(ne),1,ne^2);

VarZi = [eye(ne)            zeros(ne,ne^2)         kron(eye(ne),E_xf')   kron(E_xf',eye(ne)); ...
         zeros(ne^2,ne)     E_eps2_eps2-vecI_vecI  zeros(ne^2,ne*nx)     zeros(ne^2,ne*nx); ...
         kron(eye(ne),E_xf) zeros(ne*nx,ne^2)      kron(eye(ne),E_xf_xf) E_epsxf_xfeps; ...
         kron(E_xf,eye(ne)) zeros(ne*nx,ne^2)      E_xfeps_epsxf         kron(E_xf_xf,eye(ne))];


% Computing the variances of the state variables, using the Lyaponov routine
Var_z       = dlyap_doubling(A,B*VarZi*B',Var_zold);
Var_xf      = Var_z(1:nx,1:nx);
Var_xs      = Var_z(nx+1:2*nx,nx+1:2*nx);
Var_xf_xs   = Var_z(1:nx,nx+1:2*nx);
Var_xfxf    = Var_z(2*nx+1:2*nx+nx^2,2*nx+1:2*nx+nx^2);
Var_xs_xfxf = Var_z(nx+1:2*nx,2*nx+1:2*nx+nx^2);
Var_x       = Var_xf + Var_xs + Var_xf_xs + Var_xf_xs';

% Computing the autu-corelations 
Cov_z       = zeros(size(Var_z,1),size(Var_z,2),auto_lag);
Cov_xf      = zeros(nx,nx,auto_lag);
Cov_xs      = zeros(nx,nx,auto_lag);
Cov_x       = zeros(nx,nx,auto_lag);
for i=1:auto_lag
    Cov_z(:,:,i)  = A^i*Var_z;
    Cov_xf(:,:,i) = Cov_z(1:nx,1:nx,i);
    Cov_xs(:,:,i) = Cov_z(nx+1:2*nx,nx+1:2*nx,i);
    Cov_x(:,:,i)  = Cov_xf(:,:,i) + Cov_xs(:,:,i) + Cov_z(1:nx,nx+1:2*nx,i) + Cov_z(nx+1:2*nx,1:nx,i);
end

% Moments of y: Recall y = C*z+d
C          = [gx gx reshape(gxx,ny,nx*nx)];
d          = g0;
E_y        = C*E_z+d;
Var_y      = C*Var_z*C';
Cov_y      = zeros(ny,ny,auto_lag);
for i=1:auto_lag
    Cov_y(:,:,i) = C*Cov_z(:,:,i)*C';
end

% output
out.E_z         = E_z;
out.E_xf        = E_xf;
out.E_xs        = E_xs;
out.E_xf_xf     = E_xf_xf;
out.Var_z       = Var_z;
out.Var_xf      = Var_xf;
out.Var_xs      = Var_xs;
out.Var_xf_xs   = Var_xf_xs;
out.Var_xfxf    = Var_xfxf;
out.Var_xs_xfxf = Var_xs_xfxf;
out.Var_x       = Var_x;
out.Cov_z       = Cov_z;
out.Cov_xf      = Cov_xf;
out.Cov_xs      = Cov_xs;
out.Cov_x       = Cov_x;
out.E_y         = E_y;
out.Var_y       = Var_y;
out.Cov_y       = Cov_y;
out.C           = C;